from .wideresnet import build_wideresnet


def build_model(args):
    if args.model == "wideresnet":
        model = build_wideresnet(depth=args.depth,
                                 widen_factor=args.widen_factor,
                                 dropout=args.dropout,
                                 num_classes= args.num_classes)
        return model